import asyncio
import json
import os
import socket
import threading
import time
from multiprocessing import Queue
import numpy as np
from scipy.spatial.transform import Slerp, Rotation as R
from Bvh2LiveLink.ConvertbvhToJson import generate_json, save_json
from Bvh2LiveLink.Parsebvh import parse_bvh

import websockets



def get_bvh_files(folder):
    return [os.path.join(folder, f) for f in os.listdir(folder) if f.endswith('.bvh')]

def load_bvh_file(bvh_file_path):
    source_name = os.path.splitext(os.path.basename(bvh_file_path))[0]
    all_frames = parse_bvh(bvh_file_path)
    bvh_json = generate_json(all_frames, source_name)
    return bvh_json

def lerp(loc1, loc2, t):
    loc1 = np.array(loc1)
    loc2 = np.array(loc2)
    return (1 - t) * loc1 + t * loc2

async def echo(websocket, path):
    async for message in websocket:
        print(f"接收到消息: {message}")
        await handle_message(message, websocket)

async def handle_message(message, websocket):
    if int(json.loads(message)["action_index"]):
        file_index = int(json.loads(message)["action_index"])
        await switch_bvh_file(file_index, websocket)
    else:
        await websocket.send("无效的输入，请输入有效的BVH文件编号。")
def interpolate_frames(frame1, frame2, steps):
    interpolated_frames = []
    for step in range(steps):
        t = step / float(steps)
        interpolated_frame = []
        for bvh_source1, bvh_source2 in zip(frame1['bvh_source'], frame2['bvh_source']):
            # 四元数插值
            slerp = Slerp([0, 1], R.from_quat([bvh_source1["Rotation"], bvh_source2["Rotation"]]))
            interpolated_rot = slerp(t).as_quat().tolist()
            # 线性插值
            interpolated_loc = lerp(bvh_source1["Location"], bvh_source2["Location"], t).tolist()
            interpolated_bone = {
                "Name": bvh_source1["Name"],
                "Parent": bvh_source1["Parent"],
                "Location": interpolated_loc,
                "Rotation": interpolated_rot,
                "Scale": bvh_source1["Scale"]
            }
            interpolated_frame.append(interpolated_bone)
        interpolated_frames.append(interpolated_frame)
    interpolated_frames_json = [{"bvh_source": frame_data} for frame_data in interpolated_frames]
    return interpolated_frames_json

async def switch_bvh_file(file_index, websocket):
    global bvh_files, bvh_json, current_frame_index, last_frame, switch_event
    if 0 <= file_index < len(bvh_files):
        bvh_file_path = bvh_files[file_index]
        bvh_json_new = load_bvh_file(bvh_file_path)
        bvh_json[:] = bvh_json_new
        current_frame_index[0] = 0
        switch_event.set()
        print(f"Switched to BVH file: {bvh_file_path}")
        await websocket.send(f"已切换到BVH文件: {bvh_file_path}")
    else:
        print("Error: 输入的编号无效。继续当前动作。")
        await websocket.send("输入的编号无效。继续当前动作。")

async def send_data(ip_address, port, stop_event, switch_event, bvh_json, current_frame_index, last_frame):
    cli = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    while not stop_event.is_set():
        for i in range(current_frame_index[0], len(bvh_json)):
            if stop_event.is_set():
                return
            if switch_event.is_set():
                switch_event.clear()
                transition_frames = interpolate_frames(last_frame, bvh_json[0], transition_frame)
                for frame in transition_frames:
                    try:
                        message = json.dumps(frame).encode('utf-8')
                        cli.sendto(message, (ip_address, port))
                    except Exception as e:
                        print(f"Error sending data: {e}")
                    await asyncio.sleep(frame_interval)
                current_frame_index[0] = 0
                last_frame = bvh_json[0]["bvh_source"]
                break
            try:
                message = json.dumps(bvh_json[i]).encode('utf-8')
                cli.sendto(message, (ip_address, port))
                last_frame = bvh_json[i]
            except Exception as e:
                print(f"Error sending data: {e}")
            await asyncio.sleep(frame_interval)
            save_json(bvh_json[20], "output.json")
    cli.close()

async def main():
    global bvh_files, bvh_json, current_frame_index, last_frame, switch_event
    bvh_folder = 'Data/bvh'
    bvh_files = get_bvh_files(bvh_folder)

    if not bvh_files:
        print("Error: No BVH files found in the specified folder.")
        return

    file_index = 0

    if file_index < 0 or file_index >= len(bvh_files):
        print("Error: 输入的编号无效。")
        return

    bvh_file_path = bvh_files[file_index]
    bvh_json = load_bvh_file(bvh_file_path)

    stop_event = asyncio.Event()
    switch_event = asyncio.Event()

    # 当前帧索引
    current_frame_index = [0]

    # 存储上一帧用于插值
    last_frame = bvh_json[0]["bvh_source"]

    # 启动WebSocket服务器
    async with websockets.serve(echo, "127.0.0.1", 8000, ping_interval=None, ping_timeout=None):
        print("WebSocket服务器已启动，正在监听127.0.0.1:8000")
        # 启动数据发送任务
        send_task = asyncio.create_task(send_data(ip_address, port, stop_event, switch_event, bvh_json, current_frame_index, last_frame))
        # 等待数据发送任务完成
        await send_task

if __name__ == "__main__":
    # 帧率设置
    frame_rate = 30
    frame_interval = 1 / frame_rate

    transition_frame = 10

    # IP和端口设置
    # ip_address = "192.168.1.10"
    ip_address = "192.168.31.28"
    port = 11111
    asyncio.run(main())
